package akka.agent.application

import akka.Done
import akka.agent.domain.Exchange
import akka.agent.domain.SessionEvent
import akka.javasdk.eventsourcedentity.EventSourcedEntity
import cn.hutool.log.Log

class SessionEntity : EventSourcedEntity<SessionEntity.Session, SessionEvent>() {
  private val log by lazy { Log.get(javaClass) }

  data class Session(val messages: List<Message>, val totalTokenUsage: Int) {
    companion object {
      fun empty() = Session(listOf(), 0)
    }

    fun add(message: Message) = copy(messages = messages + message)

    fun addTokenUsage(usage: Int) = copy(totalTokenUsage = totalTokenUsage + usage)
  }

  data class Message(val content: String, val type: MessageType)

  enum class MessageType {
    AI,
    USER
  }

  override fun applyEvent(event: SessionEvent): Session =
      when (event) {
        is SessionEvent.UserMessageAdded ->
            currentState()
                .add(Message(content = event.query, type = MessageType.USER))
                .addTokenUsage(usage = event.tokensUsed)
        is SessionEvent.AiMessageAdded ->
            currentState()
                .add(Message(content = event.response, type = MessageType.AI))
                .addTokenUsage(usage = event.tokensUsed)
      }

  fun getHistory(): ReadOnlyEffect<List<Message>> {
    log.debug("Getting history from ${commandContext().entityId()}")
    return effects().reply(currentState().messages)
  }

  fun addExchange(exchange: Exchange): Effect<Done> {
    val userEvt =
        SessionEvent.UserMessageAdded(
            userId = exchange.userId,
            sessionId = exchange.sessionId,
            query = exchange.userQuestion,
            tokensUsed = exchange.queryTokensCount,
        )
    val assistantEvt =
        SessionEvent.AiMessageAdded(
            userId = exchange.userId,
            sessionId = exchange.sessionId,
            response = exchange.assistantResponse,
            tokensUsed = exchange.responseTokensCount)
    return effects().persist(userEvt, assistantEvt).thenReply { Done.getInstance() }
  }
}
